from langchain.chains.llm import LLMChain
from langchain.prompts import PromptTemplate
from langchain_community.chat_models import ChatZhipuAI

from my_intelligent_customer.utils.prompts import STAGE_ANALYZER_INCEPTION_PROMPT


class StageAnalyzerChain(LLMChain):
    """Chain to analyze which conversation stage should the conversation move into."""

    @classmethod
    def from_llm(cls, llm: ChatZhipuAI, verbose: bool = True) -> LLMChain:
        """Get the response parser."""
        stage_analyzer_inception_prompt_template = STAGE_ANALYZER_INCEPTION_PROMPT
        prompt = PromptTemplate(
            template=stage_analyzer_inception_prompt_template,
            input_variables=[
                "summary",
                "conversation_history",
                "conversation_stages",
                "conversation_stage_id",
            ],
        )
        print(f"STAGE ANALYZER PROMPT {prompt}")
        return cls(prompt=prompt, llm=llm, verbose=verbose)
